import chisel3._
import chisel3.util._

class ICacheIO(xlen:Int) extends  Bundle {
  // IFU <=> ICache
  val addr  = Input(UInt(xlen.W))
  val data  = Output(UInt(xlen.W))

  val avalid = Input(Bool())
  val aready = Output(Bool())
  val rvalid = Output(Bool())
  val rready = Input(Bool())

  // val valid = Input(Bool())
  // val ready = Output(Bool())

  // IFU <=> Master interface
  val ACLK    = Input(Clock())
  val ARESETn = Input(Reset())
  val ma      = new MessageAXI4MasterInterface(xlen)

  // IDU <=> ICache
  val inst = Input(UInt(xlen.W))
  val inst_valid = Input(Bool())
}

class ICache(xlen:Int, bsize:Int, bnum:Int) extends Module {
  val io = IO(new ICacheIO(xlen))

  val m:Int    = log2Ceil(bsize)
  val n:Int    = log2Ceil(bnum)
  val line:Int = xlen - m - n + bsize*8 + 1

  val icache   = RegInit(VecInit(Seq.fill(bnum)(0.U(line.W))))

  val inst_addr = RegInit(0.U(xlen.W))
  inst_addr := Mux(io.avalid & io.aready, io.addr, inst_addr)

  val atag  = Wire(UInt((xlen-m-n).W))
  atag := inst_addr(xlen - 1, m+n)

  val index  = Wire(UInt(n.W))
  val itag   = Wire(UInt((xlen-m-n).W))
  val ivalid = Wire(Bool())
  index  := inst_addr(m+n-1, m)
  itag   := icache(index)(line-1, (bsize*8+1))
  ivalid := icache(index)(bsize*8)

  val vec_size:Int = bsize/(xlen/8)-1
  val read_data    = Reg(Vec(vec_size, UInt(xlen.W)))
  val data_index   = Wire(UInt(m.W))

  val rvalid = RegInit(0.B)  // for sending data
  val aready = RegInit(1.B)  // for receieves address

  // synchronize with outside clock
  withClockAndReset(io.ACLK, !(io.ARESETn).asBool) {
    /* Master interface <=> Arbiter */
    // Read address channel
    val ARVALID = RegInit(0.B)
    val ARADDR  = RegInit(0.U(xlen.W))
    // Read data channel
    val RREADY  = RegInit(0.B)

    val counter = RegInit(0.U(4.W))

    val sIDEL :: sIREAD :: sMREAD :: Nil = Enum(3)
    val curr_status = RegInit(sIDEL)
    val next_status = WireDefault(sIDEL)

    // status transfer
    curr_status := next_status
    switch (curr_status) {
      is (sIDEL) {
        when (io.avalid & io.aready) {
          next_status := sIREAD
        } .otherwise {
          next_status := sIDEL
        }
      }

      // read icache
      is (sIREAD) {
        when (ivalid && atag===itag) {
          next_status := sIDEL
        } .otherwise {
          next_status := sMREAD
        }
      }

      // read mem
      is (sMREAD) {
        when (io.ma.rvalid && io.ma.rready && io.ma.rlast) {
          next_status := sIREAD
        } .otherwise {
          next_status := sMREAD
        }
      }
    }


    counter := Mux(curr_status === sIREAD && next_status === sMREAD, 0.U, Mux(io.ma.rvalid & io.ma.rready & ~io.ma.rlast, counter + 1.U, counter))

    //---------------------
    //Read address channel
    //---------------------
    ARVALID := Mux(io.ma.arvalid && io.ma.arready, 0.B, Mux(curr_status === sIREAD && next_status === sMREAD, 1.B, ARVALID))
    ARADDR  := Mux(curr_status === sIREAD && next_status === sMREAD, inst_addr & ~((bsize - 1).U(xlen.W)), ARADDR)

    //---------------------
    //Read data channel
    //---------------------
    RREADY := Mux(io.ma.rvalid && !io.ma.rready, 1.B, Mux(io.ma.rready, 0.B, RREADY))

    // ICache sends data
    rvalid := Mux(io.rvalid & io.rready, 0.B, Mux(curr_status === sIREAD, Mux(ivalid && atag===itag, 1.B, rvalid), rvalid))
    io.rvalid := rvalid

    // ICache receieves address
    aready := Mux(io.avalid & io.aready, 0.B, Mux(curr_status === sIDEL, 1.B, aready))
    io.aready := aready

    // Read address channel
    io.ma.arvalid := ARVALID
    io.ma.araddr  := ARADDR

    // Read data channel
    io.ma.rready := RREADY

    if (vec_size != 0) {
      read_data(counter) := Mux(io.ma.rvalid & io.ma.rready & ~io.ma.rlast, io.ma.rdata, read_data(counter))
    }
  }

  if (vec_size != 0) {  // block size is not 1 word
    data_index := inst_addr(m-1, 2)
  } else {
    data_index := 0.U
  }
  io.data    := (0 to vec_size).foldLeft(0.U(xlen.W)) { case (acc, i) =>
    Mux(data_index === i.U, icache(index)(i*xlen + xlen - 1, i*xlen), acc)
  }

  // Write address channel
  io.ma.awvalid := 0.B
  io.ma.awaddr  := 0.U
  io.ma.awid    := 0.U            // default is zero
  io.ma.awlen   := vec_size.U     // total number of transfers: Length = AWLEN + 1
  io.ma.awburst := 1.U            // INCR burst type
  io.ma.awsize  := Mux((xlen == 32).B, 2.U(3.W), Mux((xlen == 64).B, 3.U(3.W), (xlen/8).U))

  // Write data channel
  io.ma.wvalid := 0.B
  io.ma.wdata  := 0.U
  io.ma.wstrb  := 0.U
  io.ma.wlast  := 1.B

  // Write response channel
  io.ma.bready  := 0.U

  // Read address channel
  io.ma.arid    := 0.U          // default is zero
  io.ma.arlen   := vec_size.U   // total number of transfers: Length = AWLEN + 1
  io.ma.arburst := 1.U          // INCR burst type
  io.ma.arsize  := Mux((xlen == 32).B, 2.U(3.W), Mux((xlen == 64).B, 3.U(3.W), (xlen/8).U))

  // tag, valid, block and excute fence.i instruction
  (0 until bnum).foreach { i =>
    if (vec_size != 0) {
      icache(i) := Mux(io.inst_valid && io.inst(6, 0) === "b0001111".asUInt, 0.U, Mux(index === i.U, Mux(io.ma.rvalid && io.ma.rready, Cat(atag, 1.B, io.ma.rdata, read_data.asUInt), icache(i)), icache(i)))
    } else {
      icache(i) := Mux(io.inst_valid && io.inst(6, 0) === "b0001111".asUInt, 0.U, Mux(index === i.U, Mux(io.ma.rvalid && io.ma.rready, Cat(atag, 1.B, io.ma.rdata), icache(index)), icache(i)))
    }
  }
}
